import os
import tkinter as tk
from tkinter.filedialog import askopenfilename
import argparse
from sys import platform
# from models import *  # set ONNX_EXPORT in models.py
# from utils.datasets import *
# from utils.utils import *
# from PIL import Image, ImageTk


# import torch
# from torchvision import transforms

# from model import AlexNet
# from networks.ClassicNetwork.ResNet import ResNet50


class MainWindow:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title('基于yolo3剪枝的识别')

        self.frame_left = tk.Frame(self.root, padx=1, pady=5, bg="#aaaaaa")
        self.frame_left.pack(padx=5, pady=10, fill="y", side=tk.LEFT)

        self.frame_right = tk.Frame(self.root, padx=5, pady=5)
        self.frame_right.pack(padx=5, pady=10, fill="y", side=tk.LEFT)

        self.button_loadModel = tk.Button(self.frame_right, text='加载模型', command=self.loadModel, width=15, height=2)
        self.button_loadModel.pack(fill="x")

        self.button_loadImage = tk.Button(self.frame_right, text='加载图像', command=self.loadimg, width=15, height=2)
        self.button_loadImage.pack(fill="x")
        self.button_loadImage.config(state=tk.DISABLED)

        self.button_predict = tk.Button(self.frame_right, text='识别一下', command=self.predict, width=15, height=2)
        self.button_predict.pack(fill="x")
        self.button_predict.config(state=tk.DISABLED)

        self.label_info = tk.Label(self.frame_right, font=('宋体', 12), justify=tk.LEFT, padx=2, pady=30)
        self.label_info.pack(fill="x")
        self.label_info.config(text="请载入一个模型文件")

        self.canvas = tk.Canvas(self.frame_left, bg='gray', height=300, width=400)
        self.canvas.pack(fill='x', expand='yes')

        self.root.mainloop()

    def load_image_to_canvas(self, file_path):
        """把给定路径的图像加载入self.img 并绘制到canvas"""

        def resize(w_box, h_box, pil_image):  # 参数是：要适应的窗口宽、高、Image.open后的图片
            w, h = pil_image.size  # 获取图像的原始大小
            f1 = 1.0 * w_box / w
            f2 = 1.0 * h_box / h
            factor = min([f1, f2])
            width = int(w * factor)
            height = int(h * factor)
            return pil_image.resize((width, height), Image.ANTIALIAS)

        try:
            img = Image.open(file_path)
            self.img = img
            img_w, img_h = img.size
            if img_w > 400:
                img_w = 400
                img_h = img_h * (400 / img_w)
                img = resize(img_w, img_h, img)
            self.pil_img = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
            self.canvas.update()  # 获取宽高之前要先对于这个组件update()
            x, y = 0, (self.canvas.winfo_height() - img.size[1]) / 2
            self.canvas.create_image(x, y, anchor='nw', image=self.pil_img)
        except Exception as e:
            self.label_info.config(text="图片载入出错")
        finally:
            self.button_predict.config(state=tk.NORMAL)
            self.label_info.config(text="图片已载入\n点击预测按钮")

    def predict(self):
        with torch.no_grad():
            detect()
        photoPath = " "
        self.load_image_to_canvas(photoPath)




    def loadModel(self):
        """载入指定的模型"""
        try:
            default_dir = os.getcwd()
            modelPath = askopenfilename(title='选择一个模型文件',
                                        initialdir=(os.path.expanduser(default_dir)),
                                        filetypes=[('weights文件', '*..weights'), ('All Files', '*')])
            if modelPath == "":
                return


            # self.label_info.config(text="载入模型中……")
            # model = ResNet50(num_classes=11)
            # model.load_state_dict(torch.load(modelPath))
            # model.eval()
            # self.model = model
        except Exception as e:
            self.label_info.config(text="模型载入出错")
        finally:
            self.button_loadImage.config(state=tk.NORMAL)
            self.label_info.config(text="请打开一张图片")

    def loadimg(self):
        """载入指定的jpg图片"""
        default_dir = os.getcwd()
        photoPath = askopenfilename(title='打开一个照片（jpg格式）',
                                    initialdir=(os.path.expanduser(default_dir)),
                                    filetypes=[('jpg文件', '*.jpg'), ('All Files', '*')])
        if photoPath == "":
            return
        self.load_image_to_canvas(photoPath)







def detect(save_txt=False, save_img=False):
    img_size = (320, 192) if ONNX_EXPORT else opt.img_size  # (320, 192) or (416, 256) or (608, 352) for (height, width)
    out, source, weights, half, view_img = opt.output, opt.source, opt.weights, opt.half, opt.view_img
    webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')

    # Initialize
    device = torch_utils.select_device(device='cpu' if ONNX_EXPORT else opt.device)
    if os.path.exists(out):
        shutil.rmtree(out)  # delete output folder
    os.makedirs(out)  # make new output folder

    # Initialize model
    model = Darknet(opt.cfg, img_size)

    # Load weights
    attempt_download(weights)
    if weights.endswith('.pt'):  # pytorch format
        model.load_state_dict(torch.load(weights, map_location=device)['model'])
    else:  # darknet format
        _ = load_darknet_weights(model, weights)

    # Fuse Conv2d + BatchNorm2d layers
    # model.fuse()

    # Eval mode
    model.to(device).eval()

    # Export mode
    if ONNX_EXPORT:
        img = torch.zeros((1, 3) + img_size)  # (1, 3, 320, 192)
        torch.onnx.export(model, img, 'weights/export.onnx', verbose=True)
        return

    # Half precision
    half = half and device.type != 'cpu'  # half precision only supported on CUDA
    if half:
        model.half()

    # Set Dataloader
    vid_path, vid_writer = None, None
    if webcam:
        view_img = True
        torch.backends.cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=img_size, half=half)
    else:
        save_img = True
        dataset = LoadImages(source, img_size=img_size, half=half)

    # Get classes and colors
    classes = load_classes(parse_data_cfg(opt.data)['names'])
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(classes))]

    # Run inference
    t0 = time.time()
    for path, img, im0s, vid_cap in dataset:
        t = time.time()

        # Get detections
        img = torch.from_numpy(img).to(device)
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        pred, _ = model(img)

        if opt.half:
            pred = pred.float()

        for i, det in enumerate(non_max_suppression(pred, opt.conf_thres, opt.nms_thres)):  # detections per image
            if webcam:  # batch_size >= 1
                p, s, im0 = path[i], '%g: ' % i, im0s[i]
            else:
                p, s, im0 = path, '', im0s

            save_path = str(Path(out) / Path(p).name)
            s += '%gx%g ' % img.shape[2:]  # print string
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += '%g %ss, ' % (n, classes[int(c)])  # add to string

                # Write results
                for *xyxy, conf, _, cls in det:
                    if save_txt:  # Write to file
                        with open(save_path + '.txt', 'a') as file:
                            file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf))

                    if save_img or view_img:  # Add bbox to image
                        label = '%s %.2f' % (classes[int(cls)], conf)
                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)])

            print('%sDone. (%.3fs)' % (s, time.time() - t))

            # Stream results
            if view_img:
                cv2.imshow(p, im0)

            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'images':
                    cv2.imwrite(save_path, im0)
                else:
                    if vid_path != save_path:  # new video
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.release()  # release previous video writer

                        fps = vid_cap.get(cv2.CAP_PROP_FPS)
                        w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                        h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h))
                    vid_writer.write(im0)

    if save_txt or save_img:
        print('Results saved to %s' % os.getcwd() + os.sep + out)
        if platform == 'darwin':  # MacOS
            os.system('open ' + out + ' ' + save_path)

    print('Done. (%.3fs)' % (time.time() - t0))










if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
    parser.add_argument('--data', type=str, default='data/coco.data', help='coco.data file path')
    parser.add_argument('--weights', type=str, default='weights/yolov3.weights', help='path to weights file')
    parser.add_argument('--source', type=str, default='data/samples', help='source')  # input file/folder, 0 for webcam
    parser.add_argument('--output', type=str, default='output', help='output folder')  # output folder
    parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
    parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression')
    parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
    parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
    parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    opt = parser.parse_args()
    print(opt)

    win = MainWindow()